{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%reload_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "import matplotlib.pylab as plt\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import seaborn as sns"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Realizability analysis\n",
    "\n",
    "Analyze the realizability theory in simulation.\n",
    "To reproduce Figure 3 of the paper, you can jump to the end of the file.\n",
    "\n",
    "### 1. Check that everything works on an example setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "from angle_set import AngleSet\n",
    "from pylocus.algorithms import procrustes\n",
    "from algorithms import reconstruct_from_angles\n",
    "\n",
    "d = 2\n",
    "N = 5\n",
    "np.random.seed(51)\n",
    "angle_set = AngleSet(N=N, d=d)\n",
    "angle_set.set_points(mode='random')\n",
    "angle_set.plot_all()\n",
    "\n",
    "points = reconstruct_from_angles(angle_set.theta_tensor)\n",
    "points_fitted, *_ = procrustes(angle_set.points, points, scale=True)\n",
    "\n",
    "plt.figure()\n",
    "plt.scatter(*points_fitted.T, label='fitted')\n",
    "plt.scatter(*angle_set.points.T, label='original', marker='x')\n",
    "plt.axis('equal')\n",
    "plt.legend(loc='best')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from pylocus.basics_angles import get_theta_tensor\n",
    "from pylocus.basics import projection\n",
    "from simulation_discrepancy import get_noisy\n",
    "from angle_set import get_linear_constraints\n",
    "\n",
    "theta_noisy = get_noisy(angle_set.theta)\n",
    "\n",
    "# reconstruct raw\n",
    "theta_noisy_tensor = get_theta_tensor(theta_noisy, angle_set.corners, angle_set.N)\n",
    "\n",
    "points_noisy_unfit = reconstruct_from_angles(theta_noisy_tensor)\n",
    "points_noisy, *_ = procrustes(angle_set.points, points_noisy_unfit, scale=True)\n",
    "\n",
    "# impose linear constraints\n",
    "Afull, bfull = get_linear_constraints(angle_set)\n",
    "bfull = bfull.flatten()\n",
    "theta_linear, __, __ = projection(theta_noisy, Afull, bfull)\n",
    "theta_linear_tensor = get_theta_tensor(theta_linear, angle_set.corners, angle_set.N)\n",
    "\n",
    "points_linear_unfit = reconstruct_from_angles(theta_linear_tensor)\n",
    "points_linear, *_ = procrustes(angle_set.points, points_linear_unfit, scale=True)\n",
    "\n",
    "# plot results\n",
    "plt.figure()\n",
    "plt.scatter(*points_linear.T, label='linear')\n",
    "plt.scatter(*points_noisy.T, label='noisy')\n",
    "plt.scatter(*angle_set.points.T, label='original', marker='x')\n",
    "plt.axis('equal')\n",
    "plt.legend(loc='best')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2. Analyze which constraints are independent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from algorithms import solve_constrained_optimization, constraint_sine_multi\n",
    "\n",
    "# choose which sine constraints to include.\n",
    "if N == 4:\n",
    "    choices_sine = range(1)\n",
    "elif N == 5:\n",
    "    choices_sine = range(3)\n",
    "elif N == 6:\n",
    "    choices_sine = range(6)  # range(7)\n",
    "\n",
    "print('minimized:')\n",
    "theta_sine, success = solve_constrained_optimization(theta_noisy,\n",
    "                                                     angle_set.corners,\n",
    "                                                     Afull=Afull,\n",
    "                                                     bfull=bfull,\n",
    "                                                     N=N,\n",
    "                                                     choices_sine=choices_sine)\n",
    "[print(constraint_sine_multi(theta_sine, angle_set.corners, N, [choice])) for choice in choices_sine]\n",
    "\n",
    "print('not minimized:')\n",
    "n_available = angle_set.get_n_sine()\n",
    "other = list(range(n_available))\n",
    "[other.remove(choice) for choice in choices_sine]\n",
    "for o in other:\n",
    "    print(constraint_sine_multi(theta_sine, angle_set.corners, N, [o]))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": false
   },
   "outputs": [],
   "source": [
    "d = 2\n",
    "np.random.seed(1)\n",
    "for N in range(4, 8):\n",
    "    print('===== N={} ====='.format(N))\n",
    "    angle_set = AngleSet(N=N, d=d)\n",
    "    angle_set.set_points('random')\n",
    "    Atri, btri = angle_set.get_triangle_constraints()\n",
    "    Aray, bray = angle_set.get_ray_constraints()\n",
    "    Alinear = np.vstack((Aray, Atri))\n",
    "    blinear = np.hstack((bray, btri))\n",
    "    print('considering only the minimal constraints:')\n",
    "    print('shape:', Alinear.shape, 'rank:', np.linalg.matrix_rank(Alinear))\n",
    "\n",
    "    Apoly, bpoly = angle_set.get_polygon_constraints(range(3, N))\n",
    "    Alinear = np.vstack((Aray, Apoly))\n",
    "    blinear = np.hstack((bray, bpoly))\n",
    "    print('considering more constraints:')\n",
    "    print('shape:', Alinear.shape, 'rank:', np.linalg.matrix_rank(Alinear))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3.  Calculate discrepancy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from angle_set import create_theta\n",
    "from simulation_discrepancy import mae\n",
    "from algorithms import reconstruct_theta\n",
    "\n",
    "N = 4\n",
    "d = 2\n",
    "np.random.seed(1)\n",
    "angle_set = AngleSet(N=N, d=d)\n",
    "angle_set.set_points('random')\n",
    "\n",
    "num_sine = angle_set.get_n_sine()\n",
    "num_linear = angle_set.get_n_linear()\n",
    "choices_sine = range(num_sine)\n",
    "choices_linear = range(num_linear)\n",
    "\n",
    "# denoise\n",
    "theta_noisy = get_noisy(angle_set.theta)\n",
    "Afull, bfull = get_linear_constraints(angle_set)\n",
    "\n",
    "theta_sine, __ = solve_constrained_optimization(theta_noisy,\n",
    "                                                angle_set.corners,\n",
    "                                                Afull=Afull,\n",
    "                                                bfull=bfull,\n",
    "                                                N=N,\n",
    "                                                choices_sine=choices_sine,\n",
    "                                                choices_linear=choices_linear)\n",
    "\n",
    "# reconstruct\n",
    "theta_sine_reconstructed, points_sine = reconstruct_theta(theta_sine, angle_set.corners, angle_set.N)\n",
    "\n",
    "mae_noisy = mae(angle_set.theta, theta_noisy)\n",
    "mae_denoised = mae(angle_set.theta, theta_sine)\n",
    "mae_discrepancy = mae(theta_sine_reconstructed, theta_sine)\n",
    "print(' noisy:       {:2.2e}\\n denoised:    {:2.2e}\\n discrepancy: {:2.2e}\\n'.format(\n",
    "    mae_noisy, mae_denoised, mae_discrepancy))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create plots for paper\n",
    "\n",
    "*Below results are obtained with simulation_discrepancy.py* \n",
    "\n",
    "Use either the pre-computed results..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ylim = [1e-4, 2e-1]\n",
    "ylim_err = [1e-9, 2e-3]\n",
    "fnames = ['discrepancy_ICASSP']  # results non-learning\n",
    "plotname = 'discrepancy'\n",
    "\n",
    "fname1 = 'discrepancy_learned_ICASSP0'  # newest results for learning, before interruption at 8, i=14\n",
    "fname2 = 'discrepancy_learned_ICASSP1'  # newest results for learning, after interruption at 8, i=14\n",
    "fnames = [fname1, fname2]\n",
    "plotname = 'discrepancy_learned'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "... or newly created results."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "ylim = None\n",
    "ylim_err = None\n",
    "fnames = ['discrepancy']\n",
    "plotname = 'discrepancy'\n",
    "\n",
    "fnames = ['discrepancy_learned']\n",
    "plotname = 'discrepancy_learned'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from simulation_discrepancy import mae\n",
    "\n",
    "dfs = [pd.read_pickle('results/{}.pkl'.format(fname)) for fname in fnames]\n",
    "df = pd.concat(dfs, ignore_index=True, sort=False)\n",
    "\n",
    "df.loc[:, 'discrepancy_sine'] = df.apply(lambda row: mae(row.theta_sine, row.theta_sine_reconstructed), axis=1)\n",
    "df.loc[:, 'discrepancy_noisy'] = df.apply(lambda row: mae(row.theta_noisy, row.theta_noisy_reconstructed), axis=1)\n",
    "df.loc[:, 'sine'] = df.apply(lambda row: row.n_sine > 0, axis=1)\n",
    "df.tail()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "from scipy.special import binom\n",
    "from helpers import savefig\n",
    "\n",
    "#function used to aggregate measurements\n",
    "#function = np.median\n",
    "function = np.mean\n",
    "\n",
    "fig, axs = plt.subplots(2, 1, sharex=True, sharey=False)\n",
    "fig.set_size_inches(3, 4)\n",
    "i = 0\n",
    "Ns = sorted(df.N.unique())\n",
    "\n",
    "kwargs = dict(linewidth=0.0, s=20)\n",
    "sns.set_palette(sns.color_palette('Blues_d', n_colors=len(Ns)))\n",
    "\n",
    "for N, df_N in df.groupby('N'):\n",
    "    N = int(N)\n",
    "    df_N = df_N.loc[df_N.success.values.astype(np.bool)]\n",
    "    df_mean = df_N.groupby('n_total').aggregate(function, axis=0)\n",
    "    df_mean.reset_index(inplace=True, drop=False)\n",
    "\n",
    "    discrepancy_noise = df_mean.discrepancy_noisy.values[0]\n",
    "    error_noise = df_mean.error_noisy.values[0]\n",
    "\n",
    "    ################ DISCREPANCY PLOT ####################\n",
    "    ax = axs[0]\n",
    "    label = 'N = {}'.format(N)\n",
    "    color = 'C{}'.format(i)\n",
    "    ax.set_yscale('log')\n",
    "    sns.scatterplot(data=df_mean,\n",
    "                    x='n_total',\n",
    "                    y='discrepancy_sine',\n",
    "                    ax=ax,\n",
    "                    label=label,\n",
    "                    style='sine',\n",
    "                    color=color,\n",
    "                    **kwargs)\n",
    "    ax.set_yscale('log')\n",
    "    if ylim_err is not None:\n",
    "        ax.set_ylim(*ylim)\n",
    "    ax.set_xlabel('number of constraints')\n",
    "    ax.set_ylabel('discrepancy')\n",
    "    ax.legend().set_visible(False)\n",
    "\n",
    "    ################ ACCURACY PLOT ####################\n",
    "    ax = axs[1]\n",
    "    label = None\n",
    "    ax.set_yscale('log')\n",
    "    sns.scatterplot(data=df_mean, x='n_total', y='error_sine', ax=ax, label=label, style='sine', color=color, **kwargs)\n",
    "    ax.axis(yscale='log')\n",
    "    if ylim_err is not None:\n",
    "        ax.set_ylim(*ylim_err)\n",
    "    ax.set_xlabel('number of constraints')\n",
    "    ax.set_ylabel('accuracy')\n",
    "    ax.legend().set_visible(False)\n",
    "    i += 1\n",
    "\n",
    "################ FANCY LEGEND ####################\n",
    "# hardcoded positions of labels\n",
    "y_offset = -3.8\n",
    "y2_offset = -8.4\n",
    "xs = [-5.0, 14.0, 42.0, 85.0, 130.0]\n",
    "ys = [10**(y_offset)] * len(Ns)\n",
    "ys[0] = 10**(y_offset + 0.3)\n",
    "ys2 = [10**(y2_offset)] * len(Ns)\n",
    "ys2[0] = 10**(y2_offset + 0.7)\n",
    "for j, N in enumerate(Ns):\n",
    "    axs[0].annotate(s='N={}'.format(int(N)), xy=(xs[j], ys[j]), color='C{}'.format(j))\n",
    "    axs[1].annotate(s='N={}'.format(int(N)), xy=(xs[j], ys2[j]), color='C{}'.format(j))\n",
    "\n",
    "################ DECORATIONS ####################\n",
    "if (plotname != 'discrepancy') and (ylim is not None):\n",
    "    axs[0].get_yaxis().set_ticklabels([])\n",
    "    axs[1].get_yaxis().set_ticklabels([])\n",
    "    axs[0].set_ylabel('')\n",
    "    axs[1].set_ylabel('')\n",
    "axs[1].grid()\n",
    "axs[0].grid()\n",
    "savefig('plots/{}.pdf'.format(plotname), fig=fig)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
